#include <croco/ssh/ssh.hpp>
#include <croco/common_exceptions.hpp>
#include <croco/ip/tcp.hpp>
#include <sstream>
#include <cassert>

#include <libssh/libssh.h>
#if !defined LIBSSH_VERSION_MAJOR || !defined LIBSSH_VERSION_MINOR
    #error Unable to detect version of the libssh library
#endif
#if LIBSSH_VERSION_MAJOR != 0 || LIBSSH_VERSION_MINOR < 5
    #error Unsupported version of the libssh library. libssh in version >= 0.5 is expected
#endif
#include <libssh/server.h>

namespace csh = croco::ssh;
namespace cip = croco::ip;
namespace cros = croco::os;

namespace {

class libssh_initializer
{
public:
    libssh_initializer()
    {
        if(ssh_init() != SSH_OK)
        {
            throw croco::libssh_error() << croco::more_info("libssh initialization failed");
        }
    }
    ~libssh_initializer()
    {
        if(ssh_finalize() != SSH_OK)
        {
            std::abort();
        }
    }
};

// NOTE: libssh library initialization
static libssh_initializer libsshinit;

template<typename T, void (*release)(T*)>
class libssh_struct_raii
{
public:
    libssh_struct_raii(T* r, const std::string& err = "resource acquisition error")
    : m_r(r)
    {
        if(m_r == 0)
        {
            throw croco::libssh_error() << croco::more_info(err);
        }
    }
    ~libssh_struct_raii()
    {
        assert(m_r);
        release(m_r);
        m_r = 0;
    }
    operator T*() { return m_r; }
    operator const T*() const { return m_r; }

private:
    // uncopyable
    libssh_struct_raii(const libssh_struct_raii&);
    libssh_struct_raii& operator=(const libssh_struct_raii&);

    T* m_r;
};

typedef libssh_struct_raii<ssh_bind_struct, ssh_bind_free> libssh_bind_raii;
typedef libssh_struct_raii<ssh_session_struct, ssh_free> libssh_session_raii;

template<typename T>
void throw_on_libssh_error(int value, T* sshStruct = 0)
{
    assert(sshStruct);
    if(value != SSH_OK)
    {
        std::stringstream ss;
        ss << ssh_get_error(sshStruct);
        throw croco::libssh_error() << croco::more_info(ss.str());
    }
}

bool process_auth_request(ssh_message_struct* message, csh::integration& intgr)
{
    assert(message);
    assert(ssh_message_type(message) == SSH_REQUEST_AUTH);
    int supportedAuthMethods = 0;
    const int implAuthMethods = intgr.get_auth_methods();
    assert(implAuthMethods > 0);
    if(implAuthMethods & csh::integration::auth_no_auth)
        supportedAuthMethods |= SSH_AUTH_METHOD_NONE;
    if(implAuthMethods & csh::integration::auth_user_passd)
        supportedAuthMethods |= SSH_AUTH_METHOD_PASSWORD;
    if(implAuthMethods & csh::integration::auth_public_key)
        supportedAuthMethods |= SSH_AUTH_METHOD_PUBLICKEY;
    if(implAuthMethods & csh::integration::auth_host_based)
        supportedAuthMethods |= SSH_AUTH_METHOD_HOSTBASED;
    ssh_message_auth_set_methods(message, supportedAuthMethods);
    const int authMethod = ssh_message_subtype(message);
    bool status = false;
    switch(authMethod)
    {
        default:
            assert(!"invalid auth method");
        case SSH_AUTH_METHOD_UNKNOWN:
            status = false;
            break;
        case SSH_AUTH_METHOD_NONE:
            status = implAuthMethods & csh::integration::auth_no_auth;
            intgr.unrestricted_auth();
            break;
        case SSH_AUTH_METHOD_PASSWORD:
            status = intgr.verify_user_passwd(ssh_message_auth_user(message), ssh_message_auth_password(message));
            break;
        case SSH_AUTH_METHOD_PUBLICKEY:
            status = intgr.verify_key("");    // TODO
            break;
        case SSH_AUTH_METHOD_HOSTBASED:
            status = intgr.verify_host("");   // TODO
            break;
        case SSH_AUTH_METHOD_INTERACTIVE:
            status = false;
            break;
    }
    if(status)
    {
        ssh_message_auth_reply_success(message, 0);
    }
    return status;
}

bool process_channel_open_request(ssh_message_struct* message, ssh_channel_struct** channel)
{
    assert(message);
    assert(ssh_message_type(message) == SSH_REQUEST_CHANNEL_OPEN);
    const int channelType = ssh_message_subtype(message);
    bool status = false;
    switch(channelType)
    {
        default:
            assert(!"invalid channel type");
        case SSH_CHANNEL_UNKNOWN:
            status = false;
            break;
        case SSH_CHANNEL_SESSION:
            {
                *channel = ssh_message_channel_request_open_reply_accept(message);
                assert(*channel);
                ssh_channel_set_blocking(*channel, 0);
                status = true;
                break;
            }
        case SSH_CHANNEL_DIRECT_TCPIP:
            status = false;
            break;
        case SSH_CHANNEL_FORWARDED_TCPIP:
            status = false;
            break;
        case SSH_CHANNEL_X11:
            status = false;
            break;
    }
    if(status)
    {
        ssh_message_channel_request_reply_success(message);
    }
    return status;
}

void set_pty_size(csh::integration& intgr, ssh_message_struct* message)
{
    assert(message);
    assert(ssh_message_type(message) == SSH_REQUEST_CHANNEL);
    assert(
        ssh_message_subtype(message) == SSH_CHANNEL_REQUEST_PTY ||
        ssh_message_subtype(message) == SSH_CHANNEL_REQUEST_WINDOW_CHANGE
    );
    intgr.set_pty_size(
        csh::integration::pty_width(ssh_message_channel_request_pty_width(message)),
        csh::integration::pty_height(ssh_message_channel_request_pty_height(message)),
        csh::integration::pty_pxwidth(ssh_message_channel_request_pty_pxwidth(message)),
        csh::integration::pty_pxheight(ssh_message_channel_request_pty_pxheight(message))
    );
}

bool process_channel_request(ssh_message_struct* message, csh::integration& intgr)
{
    assert(message);
    assert(ssh_message_type(message) == SSH_REQUEST_CHANNEL);
    const int channelRequest = ssh_message_subtype(message);
    bool status = false;
    switch(channelRequest)
    {
        default:
            assert(!"invalid channel request");
        case SSH_CHANNEL_REQUEST_UNKNOWN:
            status = false;
            break;
        case SSH_CHANNEL_REQUEST_PTY:
            intgr.set_pty_name(ssh_message_channel_request_pty_term(message));
            set_pty_size(intgr, message);
            status = true;
            break;
        case SSH_CHANNEL_REQUEST_EXEC:
            status = false;
            break;
        case SSH_CHANNEL_REQUEST_SHELL:
            status = true;
            break;
        case SSH_CHANNEL_REQUEST_ENV:
            intgr.set_env_var(ssh_message_channel_request_env_name(message), ssh_message_channel_request_env_value(message));
            status = true;
            break;
        case SSH_CHANNEL_REQUEST_SUBSYSTEM:
            status = false;
            break;
        case SSH_CHANNEL_REQUEST_WINDOW_CHANGE:
            set_pty_size(intgr, message);
            status = true;
            break;
    }
    if(status)
    {
        ssh_message_channel_request_reply_success(message);
    }
    return status;
}

struct msg_callback_ctx
{
    msg_callback_ctx(csh::integration& i)
    : channel(0), intgr(i)
    {
    }

    ssh_channel_struct* channel;
    csh::integration& intgr;
    boost::exception_ptr active_exception;
};

int message_callback(ssh_session_struct* session, ssh_message_struct* message, void *data)
{
    assert(session != 0 && message != 0 && data != 0);
    msg_callback_ctx* ctx = static_cast<msg_callback_ctx*>(data);
    int type = ssh_message_type(message);
    bool responseSent = false;
    try
    {
        switch(type)
        {
            case SSH_REQUEST_AUTH:
                responseSent = process_auth_request(message, ctx->intgr);
                break;
            case SSH_REQUEST_CHANNEL_OPEN:
                responseSent = process_channel_open_request(message, &ctx->channel);
                break;
            case SSH_REQUEST_CHANNEL:
                responseSent = process_channel_request(message, ctx->intgr);
                break;
            case SSH_REQUEST_SERVICE:
            case SSH_REQUEST_GLOBAL:
            default:
                responseSent = false;
        }
    }
    catch(...)
    {
        ctx->active_exception = boost::current_exception();
    }
    return responseSent ? 0 : 1;
}

} // anonymous namespace

bool csh::integration::verify_user_passwd(const std::string& user, const std::string& passwd)
{
    BOOST_THROW_EXCEPTION(not_implemented());
}

bool csh::integration::verify_key(const std::string& publicKey)
{
    BOOST_THROW_EXCEPTION(not_implemented());
}

bool csh::integration::verify_host(const std::string& hostName)
{
    BOOST_THROW_EXCEPTION(not_implemented());
}

struct csh::interface::sshimpl
{
    sshimpl(csh::integration& intgr)
    : bind(ssh_bind_new()),
      session(ssh_new()),
      callback_ctx(intgr)
    {
    }

    libssh_bind_raii bind;
    libssh_session_raii session;
    msg_callback_ctx callback_ctx;
};

csh::interface::interface(csh::integration& intgr,
    const std::string& rsaKeyPath,
    const std::string& dsaKeyPath,
    const std::string& banner,
    cros::select_timeout authTimeout,
    bool verbose)
: m_sshimpl(new sshimpl(intgr)),
  m_rsaKeyPath(rsaKeyPath),
  m_dsaKeyPath(dsaKeyPath),
  m_banner(banner),
  m_authTimeout(authTimeout),
  m_verbosity(verbose ? SSH_LOG_PACKET : SSH_LOG_NOLOG)
{
    throw_on_libssh_error(ssh_bind_options_set(m_sshimpl->bind, SSH_BIND_OPTIONS_RSAKEY, m_rsaKeyPath.c_str()), &m_sshimpl->bind);
    throw_on_libssh_error(ssh_bind_options_set(m_sshimpl->bind, SSH_BIND_OPTIONS_DSAKEY, m_dsaKeyPath.c_str()), &m_sshimpl->bind);
    throw_on_libssh_error(ssh_bind_options_set(m_sshimpl->bind, SSH_BIND_OPTIONS_LOG_VERBOSITY, &m_verbosity), &m_sshimpl->bind);
    throw_on_libssh_error(ssh_bind_options_set(m_sshimpl->bind, SSH_BIND_OPTIONS_BANNER, m_banner.c_str()), &m_sshimpl->bind);
}

csh::interface::~interface()
{
    ssh_bind_set_fd(m_sshimpl->bind, -1); // to prevent calling close() on serverSock.native_handle() in the ssh_bind_free function
    ssh_disconnect(m_sshimpl->session);
    delete m_sshimpl;
}

void csh::interface::accept_connection(cip::server_tcp& socket)
{
    ssh_bind_set_fd(m_sshimpl->bind, socket.native_handle());
    throw_on_libssh_error(ssh_options_set(m_sshimpl->session, SSH_OPTIONS_LOG_VERBOSITY, &m_verbosity), &m_sshimpl->session);
    throw_on_libssh_error(ssh_bind_accept(m_sshimpl->bind, m_sshimpl->session), &m_sshimpl->bind);
    throw_on_libssh_error(ssh_handle_key_exchange(m_sshimpl->session), &m_sshimpl->session);
    ssh_set_message_callback(m_sshimpl->session, message_callback, &m_sshimpl->callback_ctx);
    ssh_set_blocking(m_sshimpl->session, 0);
}

bool csh::interface::authenticate(const std::vector<int>& terms)
{
    while(m_sshimpl->callback_ctx.channel == 0)
    {
        process_protocol_msgs();
        if(!ssh_is_connected(m_sshimpl->session))
        {
            // connection lost
            return false;
        }
        cros::select_fds sets;
        sets.add_read(get_session_fd());
        for(std::size_t i = 0; i < terms.size(); i++)
        {
            sets.add_interrupt(terms[i]);
        }
        std::set<cros::native_handle_type> signaled = cros::do_select(sets, m_authTimeout);
        if(sets.interrupted(signaled))
        {
            return false;
        }
        if(signaled.empty())
        {
            // waiting for auth data timed-out
            return false;
        }
    }
    if(m_sshimpl->callback_ctx.channel != 0)
    {
        // NOTE: currently sending banner does not work thus is sent "manually"
        send(m_banner.c_str(), m_banner.size(), terms);
    }
    return m_sshimpl->callback_ctx.channel != 0;
}

int csh::interface::send(const char* buf, std::size_t size, const std::vector<int>& terms)
{
    assert(buf);
    assert(m_sshimpl->callback_ctx.channel);
    process_protocol_msgs();

    std::size_t sent = 0;
    while(sent < size)
    {
        if(!ssh_is_connected(m_sshimpl->session))
        {
            // connection lost
            return -1;
        }
        cros::select_fds sets;
        sets.add_write(get_session_fd());
        for(std::size_t i = 0; i < terms.size(); i++)
        {
            sets.add_interrupt(terms[i]);
        }
        std::set<cros::native_handle_type> signaled = cros::do_select(sets);
        if(signaled.empty() || sets.interrupted(signaled))
        {
            return sent;
        }
        const int lastSent = ssh_channel_write(m_sshimpl->callback_ctx.channel, buf + sent, size - sent);
        if(lastSent <= 0)
        {
            return -1;
        }
        sent += lastSent;
    }
    return sent;
}

int csh::interface::receive(char* buf, std::size_t size, const std::vector<int>& terms)
{
    assert(buf);
    assert(m_sshimpl->callback_ctx.channel);
    process_protocol_msgs();
    if(!ssh_is_connected(m_sshimpl->session))
    {
        // connection lost
        return -1;
    }
    const int availableDataNum = ssh_channel_poll(m_sshimpl->callback_ctx.channel, 0);
    if(availableDataNum < 0)
    {
        return -1;
    }
    if(availableDataNum == 0)
    {
        cros::select_fds sets;
        sets.add_read(get_session_fd());
        for(std::size_t i = 0; i < terms.size(); i++)
        {
            sets.add_interrupt(terms[i]);
        }
        std::set<cros::native_handle_type> signaled = cros::do_select(sets);
        if(sets.interrupted(signaled))
        {
            return -1;
        }
    }
    return ssh_channel_read_nonblocking(m_sshimpl->callback_ctx.channel, buf, size, 0);
}

int csh::interface::get_session_fd()
{
    return ssh_get_fd(m_sshimpl->session);
}

bool csh::interface::is_eof() const
{
    return m_sshimpl->callback_ctx.channel == 0 ? false :
        (ssh_channel_is_eof(m_sshimpl->callback_ctx.channel) != 0 ||
         ssh_channel_is_closed(m_sshimpl->callback_ctx.channel) != 0);
}

void csh::interface::process_protocol_msgs()
{
    throw_on_libssh_error(ssh_execute_message_callbacks(m_sshimpl->session), &m_sshimpl->session);
    if(m_sshimpl->callback_ctx.active_exception)
    {
        boost::rethrow_exception(m_sshimpl->callback_ctx.active_exception);
    }
}
